#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.time import Time
from geometry_msgs.msg import Twist, Point
from std_msgs.msg import String
from nav_msgs.msg import Odometry
from ball_collector_interfaces.msg import BallInfo, RobotState as RobotStateMsg
from ball_collector_interfaces.srv import SetGripper, ChangeColor
import time
import math

class StrategyTestNode(Node):
    def __init__(self):
        super().__init__('strategy_test_node')
        
        # 状态追踪
        self.current_state = "IDLE"
        self.has_ball = False
        self.target_color = "red"
        
        # 测试控制变量
        self.test_step = 0
        self.step_start_time = time.time()
        self.last_action_time = time.time()
        self.last_log_times = {}  # 用于手动实现日志节流
        
        # 创建订阅者
        self.cmd_vel_sub = self.create_subscription(
            Twist,
            'cmd_vel',
            self.cmd_vel_callback,
            10
        )
        
        self.robot_state_sub = self.create_subscription(
            RobotStateMsg,
            'robot_state',
            self.robot_state_callback,
            10
        )
        
        # 创建发布者
        self.ball_pub = self.create_publisher(
            BallInfo,
            'detected_ball',
            10
        )
        
        self.odom_pub = self.create_publisher(
            Odometry,
            'odom',
            10
        )
        
        # 创建服务
        self.gripper_service = self.create_service(
            SetGripper,
            '/hardware_interface/set_gripper',
            self.gripper_callback
        )
        
        self.color_service = self.create_service(
            ChangeColor,
            'change_target_color',
            self.color_callback
        )
        
        # 创建测试定时器
        self.test_timer = self.create_timer(1, self.run_test_sequence)
        
        # # 创建状态打印定时器
        # self.status_timer = self.create_timer(2, self.print_status)
        
        self.get_logger().info('策略测试节点已启动')
    
    def log_throttle(self, key, interval, message):
        """实现日志节流功能"""
        current_time = time.time()
        last_time = self.last_log_times.get(key, 0)
        
        if current_time - last_time >= interval:
            self.get_logger().info(message)
            self.last_log_times[key] = current_time
######################################################################
    def cmd_vel_callback(self, msg):
        """速度命令回调"""
        # 只记录有意义的变化
        if abs(msg.linear.x) > 0.01 or abs(msg.angular.z) > 0.01:
            self.get_logger().info(f"速度命令: 线速度={msg.linear.x:.2f}, 角速度={msg.angular.z:.2f}")
    
    # def robot_state_callback(self, msg):
    #     """机器人状态回调"""
    #     if msg.current_state != self.current_state or msg.has_ball != self.has_ball:
    #         self.get_logger().info(f"状态变化: {self.current_state} -> {msg.current_state}, 持球: {msg.has_ball}")
    #         self.current_state = msg.current_state
    #         self.has_ball = msg.has_ball
    #         self.target_color = msg.target_color
    def robot_state_callback(self, msg):
        """机器人状态回调"""
        # 打印收到的每条消息
        self.get_logger().debug(f"收到状态消息: state={msg.current_state}, has_ball={msg.has_ball}")
        
        if msg.current_state != self.current_state or msg.has_ball != self.has_ball:
            self.get_logger().info(f"状态变化: {self.current_state} -> {msg.current_state}, 持球: {msg.has_ball}")
            self.current_state = msg.current_state
            self.has_ball = msg.has_ball
            self.target_color = msg.target_color
        
        # 根据当前状态直接发布相应测试数据
        if msg.current_state == "SEARCHING":
            # 在搜索状态下发布球信息
            self.publish_mock_ball(True, 320, 240, 80)
            self.get_logger().info("检测到搜索状态，发布球信息")
        elif msg.current_state == "APPROACHING":
            # 在接近状态下发布更大的球
            self.publish_mock_ball(True, 320, 240, 150)
            self.get_logger().info("检测到接近状态，发布更大球信息")
#############################3#######################################################
    def gripper_callback(self, request, response):
        """夹爪服务回调"""
        self.get_logger().info(f"夹爪控制: 抓取={request.grip}, 力度={request.force}")
        response.success = True
        response.message = f"夹爪{'关闭' if request.grip else '打开'}成功"
        return response
    
    def color_callback(self, request, response):
        """颜色服务回调"""
        self.get_logger().info(f"颜色更改: {request.new_color}")
        self.target_color = request.new_color
        response.success = True
        response.current_color = request.new_color
        return response
######################################################################################   
    def publish_mock_ball(self, is_visible=True, x=320, y=240, diameter=200):
        """发布模拟球体"""
        if not is_visible:
            return
            
        msg = BallInfo()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "camera_link"
        msg.id = 1
        msg.color = self.target_color
        
        msg.center = Point()
        msg.center.x = float(x)
        msg.center.y = float(y)
        msg.center.z = 0.0
        
        msg.diameter = float(diameter)
        msg.confidence = 0.95
        msg.is_target = True
        
        self.ball_pub.publish(msg)
    
    def publish_mock_odometry(self, x=0.0, y=0.0, theta=0.0):
        """发布模拟里程计数据"""
        msg = Odometry()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "odom"
        msg.child_frame_id = "base_link"
        
        # 位置
        msg.pose.pose.position.x = x
        msg.pose.pose.position.y = y
        msg.pose.pose.position.z = 0.0
        
        # 将偏航角转换为四元数
        msg.pose.pose.orientation.w = math.cos(theta / 2.0)
        msg.pose.pose.orientation.z = math.sin(theta / 2.0)
        
        self.odom_pub.publish(msg)
######################################################################################
    def run_test_sequence(self):
        """运行测试序列"""
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        # 每次都发布里程计数据保持连接
        self.publish_mock_odometry(0.0, 0.0, 0.0)
        
        # 基于测试步骤执行不同测试
        if self.test_step == 0:
            # 初始等待测试 - 等待策略节点初始化
            if elapsed > 3.0:
                self.get_logger().info("测试步骤 1: 等待策略节点启动 - 完成")
                self.print_status()
                self.test_step = 1
                self.step_start_time = current_time
        
        elif self.test_step == 1:
            # 测试搜索状态 - 不发球信息让机器人持续搜索
            self.log_throttle("search_test", 5.0, "测试步骤 2: 搜索状态测试中...")
            
            # 不发布球，让策略节点保持搜索
            if elapsed > 5.0:
                self.get_logger().info("测试步骤 2: 搜索状态测试 - 完成")
                self.print_status()
                self.test_step = 2
                self.step_start_time = current_time
        
        elif self.test_step == 2:
            # 测试接近状态 - 发布远处的球信息
            self.log_throttle("approach_test", 5.0, "测试步骤 3: 接近状态测试中...")
            
            # 发布球信息 - 直径从小到大模拟靠近
            ball_diameter = 180 + elapsed * 10  # 直径随时间增长
            self.publish_mock_ball(True, 320, 240, ball_diameter)
            
            if elapsed > 8.0 or self.current_state == "ALIGNING":
                self.get_logger().info("测试步骤 3: 接近状态测试 - 完成")
                self.print_status()
                self.test_step = 3
                self.step_start_time = current_time
        
        elif self.test_step == 3:
            # 测试对准状态 - 发布未对准的球，然后逐渐对准
            self.log_throttle("align_test", 5.0, "测试步骤 4: 对准状态测试中...")
            
            # 球从右侧逐渐移动到中心
            x_pos = max(320, 420 - elapsed * 20)
            self.publish_mock_ball(True, x_pos, 240, 150)
            
            if elapsed > 8.0 or self.current_state == "GRASPING":
                self.get_logger().info("测试步骤 4: 对准状态测试 - 完成")
                self.print_status()
                self.test_step = 4
                self.step_start_time = current_time
        
        elif self.test_step == 4:
            # 测试抓取状态 - 发布中心对准的球
            self.log_throttle("grasp_test", 5.0, "测试步骤 5: 抓取状态测试中...")
            
            # 发布中心对准的球
            self.publish_mock_ball(True, 320, 240, 180)
            
            if elapsed > 5.0 or self.current_state == "TRANSPORTING":
                self.get_logger().info("测试步骤 5: 抓取状态测试 - 完成")
                self.print_status()
                self.test_step = 5
                self.step_start_time = current_time
        
        elif self.test_step == 5:
            # 测试运输状态 - 机器人已有球，不发球信息
            self.log_throttle("transport_test", 5.0, "测试步骤 6: 运输状态测试中...")
            
            # 模拟移动到目标区域
            x = min(2.0, elapsed * 0.2)
            y = min(1.0, elapsed * 0.1)
            self.publish_mock_odometry(x, y, 0.0)
            
            if elapsed > 10.0 or self.current_state == "RELEASING":
                self.get_logger().info("测试步骤 6: 运输状态测试 - 完成")
                self.print_status()
                self.test_step = 6
                self.step_start_time = current_time
        
        elif self.test_step == 6:
            # 测试释放状态
            self.log_throttle("release_test", 5.0, "测试步骤 7: 释放状态测试中...")
            
            # 保持在目标位置
            self.publish_mock_odometry(2.0, 1.0, 0.0)
            
            if elapsed > 5.0 or self.current_state == "RETURNING":
                self.get_logger().info("测试步骤 7: 释放状态测试 - 完成")
                self.print_status()
                self.test_step = 7
                self.step_start_time = current_time
        
        elif self.test_step == 7:
            # 测试返回状态
            self.log_throttle("return_test", 5.0, "测试步骤 8: 返回状态测试中...")
            
            # 模拟返回
            x = max(0.0, 2.0 - elapsed * 0.3)
            y = max(0.0, 1.0 - elapsed * 0.2)
            self.publish_mock_odometry(x, y, elapsed * 0.1)
            
            if elapsed > 8.0 or (self.current_state == "SEARCHING" and elapsed > 3.0):
                self.get_logger().info("测试步骤 8: 返回状态测试 - 完成")
                self.print_status()
                self.test_step = 8
                self.step_start_time = current_time
        
        elif self.test_step == 8:
            # 测试球丢失行为
            self.log_throttle("ball_lost_test", 5.0, "测试步骤 9: 球丢失行为测试中...")
            
            # 间歇性发布球信息
            ball_visible = (int(elapsed * 2) % 4 < 1)  # 25%的时间可见
            if ball_visible:
                self.publish_mock_ball(True, 320, 240, 100)
            
            if elapsed > 15.0:
                self.get_logger().info("测试步骤 9: 球丢失行为测试 - 完成")
                self.print_status()
                self.test_step = 9
                self.step_start_time = current_time
        
        elif self.test_step == 9:
            # 测试颜色切换
            self.log_throttle("color_test", 5.0, "测试步骤 10: 颜色切换测试中...")
            
            # 不同颜色的球
            colors = ["red", "green", "blue", "yellow"]
            color_index = int(elapsed) % len(colors)
            color = colors[color_index]
            
            # 每3秒切换一次颜色的球
            if current_time - self.last_action_time > 3.0:
                self.publish_mock_ball(True, 320, 240, 100)
                self.last_action_time = current_time
                
                # 切换到新颜色
                self.get_logger().info(f"尝试切换到颜色: {color}")
            
            if elapsed > 20.0:
                self.get_logger().info("测试步骤 10: 颜色切换测试 - 完成")
                self.print_status()
                self.test_step = 10
                self.step_start_time = current_time
        
        elif self.test_step == 10:
            # 测试完整收集周期
            self.log_throttle("cycle_test", 5.0, "测试步骤 11: 完整收集周期测试中...")
            
            # 根据当前状态执行相应操作
            if self.current_state == "SEARCHING":
                # 发布球信息触发接近
                self.publish_mock_ball(True, 320, 240, 80)
            
            elif self.current_state == "APPROACHING":
                # 发布逐渐增大的球模拟接近
                ball_diameter = 80 + (current_time - self.step_start_time) % 10 * 10
                self.publish_mock_ball(True, 320, 240, ball_diameter)
            
            elif self.current_state == "ALIGNING":
                # 确保球在中心
                self.publish_mock_ball(True, 320, 240, 150)
            
            elif self.current_state == "GRASPING":
                # 确保球可抓取
                self.publish_mock_ball(True, 320, 240, 180)
            
            elif self.current_state == "TRANSPORTING":
                # 模拟移动到目标区域
                elapsed_in_state = current_time - self.last_action_time
                x = min(2.0, elapsed_in_state * 0.2)
                y = min(1.0, elapsed_in_state * 0.1)
                self.publish_mock_odometry(x, y, 0.0)
            
            elif self.current_state == "RELEASING":
                # 在目标位置
                self.publish_mock_odometry(2.0, 1.0, 0.0)
            
            elif self.current_state == "RETURNING":
                # 返回起始位置
                elapsed_in_state = current_time - self.last_action_time
                x = max(0.0, 2.0 - elapsed_in_state * 0.3)
                y = max(0.0, 1.0 - elapsed_in_state * 0.2)
                self.publish_mock_odometry(x, y, elapsed_in_state * 0.1)
            
            # 跟踪状态变化
            if self.current_state != getattr(self, 'last_recorded_state', None):
                self.last_action_time = current_time
                self.last_recorded_state = self.current_state
            
            # 如果已经完成两个完整周期或者时间超过60秒，结束测试
            if elapsed > 60.0:
                self.get_logger().info("测试步骤 11: 完整收集周期测试 - 完成")
                self.print_status()
                self.test_step = 0  # 重新开始测试循环
                self.step_start_time = current_time
                self.get_logger().info("所有测试完成，重新开始测试循环")
    
    def print_status(self):
        """打印当前状态"""
        self.get_logger().info(f"当前测试步骤: {self.test_step+1}, 机器人状态: {self.current_state}, 持球: {self.has_ball}")
######################################################################################
def main(args=None):
    rclpy.init(args=args)
    node = StrategyTestNode()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        # 发布停止命令确保安全
        try:
            msg = Twist()
            msg.linear.x = 0.0
            msg.angular.z = 0.0
            node.cmd_vel_pub.publish(msg)
        except:
            pass
        
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()